#include "isp_pipeline.h"
#include <fstream>
#include <sstream>

ISP_Pipeline::ISP_Pipeline(const std::string& config_path) {
    loadConfig(config_path);
}

void ISP_Pipeline::loadConfig(const std::string& config_path) {
    std::ifstream file(config_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open config file.");
    }

    std::string line;
    while (std::getline(file, line)) {
        size_t equal_pos = line.find('=');
        if (equal_pos != std::string::npos && !line.empty() && line[0] != '#') {
            std::string key = line.substr(0, equal_pos);
            std::string value = line.substr(equal_pos + 1);

            config[key] = value;
        }
    }

    // 从配置中读取参数
    input_path = config["input_raw_image_path"];
    output_path = config["output_image_path"];
    width = std::stoi(config["width"]);
    height = std::stoi(config["height"]);
    cfa_pattern = config["cfa_pattern"];
    white_balance_r = std::stod(config["white_balance_r"]);
    white_balance_b = std::stod(config["white_balance_b"]);
    demosaicing_method = config["demosaicing_method"];
    gamma = std::stod(config["gamma"]);
    blc_value = std::stoi(config["blc_value"]);
    dpc_threshold = std::stoi(config["dpc_threshold"]);
    lsc_intensity = std::stod(config["lsc_intensity"]);
}

void ISP_Pipeline::process() {
    // 1. 读取 RAW 图像数据
    std::ifstream raw_file(input_path, std::ios::binary);
    if (!raw_file) {
        throw std::runtime_error("Failed to open RAW file.");
    }

    raw_file.seekg(0, std::ios::end);
    std::streampos fileSize = raw_file.tellg();
    raw_file.seekg(0, std::ios::beg);

    if (fileSize != width * height) {
        std::cerr << "File size mismatch. Expected: " << width * height
            << " bytes, Actual: " << fileSize << " bytes." << std::endl;
        throw std::runtime_error("File size does not match the specified width and height.");
    }

    cv::Mat raw_image(height, width, CV_8UC1);
    raw_file.read(reinterpret_cast<char*>(raw_image.data), width * height);
    raw_file.close();
    cv::imwrite("raw.png", raw_image);

    // 2. RAW 域处理
    blackLevelCalibration(raw_image);
    cv::imwrite("after_blc.png", raw_image); // 保存黑电平校准后的图像
    dynamicPixelCorrection(raw_image); // 坏点校准
    cv::imwrite("after_dpc.png", raw_image); // 保存坏点校准后的图像
    rawNoiseReduction(raw_image);
    cv::imwrite("after_rns.png", raw_image); // 保存降噪后的图像
    //lensShadingCorrection(raw_image);  // 镜头阴影校准
    //cv::imwrite("after_lsc.png", raw_image); // 保存镜头阴影校准后的图像

    // 3. Demosaicing (去马赛克)
    demosaicing(raw_image);
    cv::imwrite("after_demosaic.png", processed_image); // 保存去马赛克后的图像

    // 4. 白平衡调整
    whiteBalance(processed_image);
    cv::imwrite("after_wb.png", processed_image); // 保存白平衡调整后的图像

    // 5. 伽马校正
    gammaCorrection(processed_image);
    cv::imwrite("after_gamma.png", processed_image); // 保存伽马校正后的图像

    // 6. RGB 转 YUV
    rgbToYuv(processed_image);
    cv::imwrite("after_rgb2yuv.png", processed_image); // 保存 RGB 转 YUV 后的图像

    // 7. YUV 域处理
    contrastEnhancement(processed_image);
    cv::imwrite("after_contrast.png", processed_image); // 保存对比度增强后的图像
    sharpen(processed_image);
    cv::imwrite("after_sharpen.png", processed_image); // 保存锐化后的图像

    // 8. 颜色空间转换回 RGB（如果需要输出 RGB 格式）
    cv::cvtColor(processed_image, processed_image, cv::COLOR_YUV2BGR);
}

void ISP_Pipeline::saveOutput(const std::string& output_path) {
    cv::imwrite(output_path, processed_image);
}

// 黑电平校准
void ISP_Pipeline::blackLevelCalibration(cv::Mat& raw_image) {
    raw_image = cv::max(raw_image - blc_value, 0);
}

// RAW 域降噪
void ISP_Pipeline::rawNoiseReduction(cv::Mat& raw_image) {
    cv::GaussianBlur(raw_image, raw_image, cv::Size(3, 3), 0);
}

// 动态坏点校准
void ISP_Pipeline::dynamicPixelCorrection(cv::Mat& raw_image) {
    cv::Mat corrected = raw_image.clone();

    // 遍历图像，检测并校准坏点
    for (int y = 1; y < raw_image.rows - 1; y++) {
        for (int x = 1; x < raw_image.cols - 1; x++) {
            uchar center = raw_image.at<uchar>(y, x);
            uchar left = raw_image.at<uchar>(y, x - 1);
            uchar right = raw_image.at<uchar>(y, x + 1);
            uchar top = raw_image.at<uchar>(y - 1, x);
            uchar bottom = raw_image.at<uchar>(y + 1, x);

            // 计算周围像素的平均值
            uchar average = (left + right + top + bottom) / 4;

            // 如果当前像素与周围像素差异过大，则认为是坏点
            if (std::abs(center - average) > dpc_threshold) {
                corrected.at<uchar>(y, x) = average;
            }
        }
    }

    raw_image = corrected;
}

// 镜头阴影校准
void ISP_Pipeline::lensShadingCorrection(cv::Mat& raw_image) {
    // 创建一个镜头阴影校准掩膜
    cv::Mat mask = cv::Mat::ones(raw_image.size(), CV_32F);
    int center_x = raw_image.cols / 2;
    int center_y = raw_image.rows / 2;

    // 从中心向外减小亮度
    for (int y = 0; y < raw_image.rows; y++) {
        for (int x = 0; x < raw_image.cols; x++) {
            double distance = std::sqrt(std::pow(x - center_x, 2) + std::pow(y - center_y, 2));
            double max_distance = std::sqrt(std::pow(center_x, 2) + std::pow(center_y, 2));
            double factor = std::max(1.0 - (distance / max_distance) * (1.0 - lsc_intensity), lsc_intensity);
            mask.at<float>(y, x) = factor;
        }
    }

    // 应用掩膜
    cv::Mat float_image;
    raw_image.convertTo(float_image, CV_32F);
    cv::multiply(float_image, mask, float_image);
    float_image.convertTo(raw_image, CV_8UC1, 1.0 / 255.0, 0);
}

// Demosaicing (去马赛克)
void ISP_Pipeline::demosaicing(cv::Mat& raw_image) {
    if (demosaicing_method == "bilinear") {
        // 创建用于存储去马赛克结果的图像
        processed_image = cv::Mat(height, width, CV_8UC3, cv::Scalar(0));

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                // 根据 Bayer 滤波阵列模式 (BG) 确定当前像素的类型
                bool isGreenRow = (y % 2 == 0);
                bool isGreenCol = (x % 2 == 0);

                if (isGreenRow && isGreenCol) {
                    // G 像素（在 BG 模式中，G 位于偶数行偶数列）
                    processed_image.at<cv::Vec3b>(y, x)[1] = raw_image.at<uchar>(y, x);

                    // 计算 R 和 B 值
                    int r_sum = 0, b_sum = 0;
                    int r_count = 0, b_count = 0;

                    // 获取右下对角线像素（B）
                    if (y < height - 1 && x < width - 1) {
                        b_sum += raw_image.at<uchar>(y + 1, x + 1);
                        b_count++;
                    }
                    // 获取左下对角线像素（B）
                    if (y < height - 1 && x > 0) {
                        b_sum += raw_image.at<uchar>(y + 1, x - 1);
                        b_count++;
                    }
                    // 获取右上对角线像素（B）
                    if (y > 0 && x < width - 1) {
                        b_sum += raw_image.at<uchar>(y - 1, x + 1);
                        b_count++;
                    }
                    // 获取左上对角线像素（B）
                    if (y > 0 && x > 0) {
                        b_sum += raw_image.at<uchar>(y - 1, x - 1);
                        b_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[0] = b_sum / (b_count + 1); // B 通道

                    // 获取水平和垂直相邻像素（R）
                    if (x > 0) {
                        r_sum += raw_image.at<uchar>(y, x - 1);
                        r_count++;
                    }
                    if (x < width - 1) {
                        r_sum += raw_image.at<uchar>(y, x + 1);
                        r_count++;
                    }
                    if (y > 0) {
                        r_sum += raw_image.at<uchar>(y - 1, x);
                        r_count++;
                    }
                    if (y < height - 1) {
                        r_sum += raw_image.at<uchar>(y + 1, x);
                        r_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[2] = r_sum / (r_count + 1); // R 通道
                }
                else if (isGreenRow && !isGreenCol) {
                    // G 像素（在 BG 模式中，G 位于偶数行奇数列）
                    processed_image.at<cv::Vec3b>(y, x)[1] = raw_image.at<uchar>(y, x);

                    // 计算 R 和 B 值
                    int r_sum = 0, b_sum = 0;
                    int r_count = 0, b_count = 0;

                    // 获取水平相邻像素（B）
                    if (x > 0) {
                        b_sum += raw_image.at<uchar>(y, x - 1);
                        b_count++;
                    }
                    if (x < width - 1) {
                        b_sum += raw_image.at<uchar>(y, x + 1);
                        b_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[0] = b_sum / (b_count + 1); // B 通道

                    // 获取垂直相邻像素（R）
                    if (y > 0) {
                        r_sum += raw_image.at<uchar>(y - 1, x);
                        r_count++;
                    }
                    if (y < height - 1) {
                        r_sum += raw_image.at<uchar>(y + 1, x);
                        r_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[2] = r_sum / (r_count + 1); // R 通道
                }
                else if (!isGreenRow && isGreenCol) {
                    // G 像素（在 BG 模式中，G 位于奇数行偶数列）
                    processed_image.at<cv::Vec3b>(y, x)[1] = raw_image.at<uchar>(y, x);

                    // 计算 R 和 B 值
                    int r_sum = 0, b_sum = 0;
                    int r_count = 0, b_count = 0;

                    // 获取垂直相邻像素（B）
                    if (y > 0) {
                        b_sum += raw_image.at<uchar>(y - 1, x);
                        b_count++;
                    }
                    if (y < height - 1) {
                        b_sum += raw_image.at<uchar>(y + 1, x);
                        b_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[0] = b_sum / (b_count + 1); // B 通道

                    // 获取水平相邻像素（R）
                    if (x > 0) {
                        r_sum += raw_image.at<uchar>(y, x - 1);
                        r_count++;
                    }
                    if (x < width - 1) {
                        r_sum += raw_image.at<uchar>(y, x + 1);
                        r_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[2] = r_sum / (r_count + 1); // R 通道
                }
                else {
                    // B 像素（在 BG 模式中，B 位于奇数行奇数列）
                    processed_image.at<cv::Vec3b>(y, x)[0] = raw_image.at<uchar>(y, x);

                    // 计算 G 和 R 值
                    int g_sum = 0, r_sum = 0;
                    int g_count = 0, r_count = 0;

                    // 获取对角线相邻像素（G）
                    if (y > 0 && x > 0) {
                        g_sum += raw_image.at<uchar>(y - 1, x - 1);
                        g_count++;
                    }
                    if (y > 0 && x < width - 1) {
                        g_sum += raw_image.at<uchar>(y - 1, x + 1);
                        g_count++;
                    }
                    if (y < height - 1 && x > 0) {
                        g_sum += raw_image.at<uchar>(y + 1, x - 1);
                        g_count++;
                    }
                    if (y < height - 1 && x < width - 1) {
                        g_sum += raw_image.at<uchar>(y + 1, x + 1);
                        g_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[1] = g_sum / (g_count + 1); // G 通道

                    // 获取水平和垂直相邻像素（R）
                    if (x > 0) {
                        r_sum += raw_image.at<uchar>(y, x - 1);
                        r_count++;
                    }
                    if (x < width - 1) {
                        r_sum += raw_image.at<uchar>(y, x + 1);
                        r_count++;
                    }
                    if (y > 0) {
                        r_sum += raw_image.at<uchar>(y - 1, x);
                        r_count++;
                    }
                    if (y < height - 1) {
                        r_sum += raw_image.at<uchar>(y + 1, x);
                        r_count++;
                    }
                    processed_image.at<cv::Vec3b>(y, x)[2] = r_sum / (r_count + 1); // R 通道
                }
            }
        }
    }
    else if (demosaicing_method == "edge_aware") {
        // 使用 OpenCV 的边缘保持去马赛克算法
        cv::cvtColor(raw_image, processed_image, cv::COLOR_BayerBG2BGR_EA);
    }
    else {
        throw std::runtime_error("Unsupported demosaicing method.");
    }
}

// 白平衡调整
void ISP_Pipeline::whiteBalance(cv::Mat& image) {
    // 分离通道
    std::vector<cv::Mat> channels;
    cv::split(image, channels);

    // 调整白平衡
    channels[2] = channels[2] * static_cast<float>(white_balance_r); // R 通道
    channels[0] = channels[0] * static_cast<float>(white_balance_b); // B 通道

    // 合并通道
    cv::merge(channels, image);
}

// 伽马校正
void ISP_Pipeline::gammaCorrection(cv::Mat& image) {
    cv::Mat lookUpTable(1, 256, CV_8UC1);
    uchar* p = lookUpTable.ptr();
    for (int i = 0; i < 256; i++) {
        p[i] = cv::saturate_cast<uchar>(255.0 * pow(i / 255.0, 1.0 / gamma));
    }
    cv::LUT(image, lookUpTable, image);
}

// RGB 转 YUV
void ISP_Pipeline::rgbToYuv(cv::Mat& image) {
    cv::cvtColor(image, image, cv::COLOR_BGR2YUV);
}

// 对比度增强
void ISP_Pipeline::contrastEnhancement(cv::Mat& image) {
    // 使用 CLAHE 进行对比度增强
    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
    std::vector<cv::Mat> yuv_channels;
    cv::split(image, yuv_channels);
    clahe->apply(yuv_channels[0], yuv_channels[0]);
    cv::merge(yuv_channels, image);
}

// 锐化
void ISP_Pipeline::sharpen(cv::Mat& image) {
    cv::Mat kernel = (cv::Mat_<float>(3, 3) <<
        0, -1, 0,
        -1, 5, -1,
        0, -1, 0);
    cv::filter2D(image, image, -1, kernel);
}